import torch
from torch import nn
import torch.nn.functional as F
from torch.nn import AdaptiveAvgPool2d

#条带池化
"""
场景解析（semantic segmentation）任务中更有效地捕捉长距离依赖关系。
传统的空间池化方法通常使用固定大小的方形窗口，难以捕捉现实场景中广泛存在的各向异性上下文
"""
class StripPooling(nn.Module):
    def __init__(self, in_channels, pool_size, norm_layer, up_kwargs):
        super(StripPooling, self).__init__()
        self.pool1 = nn.AdaptiveAvgPool2d(pool_size[0])
        self.pool2 = nn.AdaptiveAvgPool2d(pool_size[1])
        self.pool3 = nn.AdaptiveAvgPool2d((1, None)) #对输入特征图在水平方向进行全局池化，保持垂直方向的尺寸不变。
        self.pool4 = nn.AdaptiveAvgPool2d((None, 1)) #对输入特征图在垂直方向进行全局池化，保持水平方向的尺寸不变。

        inter_channels=int(in_channels/4) #将输入通道数减少为原来的四分之一，用于后续的卷积操作，减少计算量。


        #self.conv1_1 和 self.conv1_2：对输入特征图进行 1x1 卷积，减少通道数，并应用归一化和激活函数。
        self.conv1_1 = nn.Sequential(nn.Conv2d(in_channels, inter_channels, 1, bias=False),
                                     norm_layer(inter_channels),
                                     nn.ReLU(True))
        self.conv1_2 = nn.Sequential(nn.Conv2d(in_channels, inter_channels, 1, bias=False),
                                     norm_layer(inter_channels),
                                     nn.ReLU(True))
        #self.conv2_0：标准的 3x3 卷积，用于提取局部特征。
        self.conv2_0 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, 1, 1, bias=False),
                                     norm_layer(inter_channels))

        self.conv2_1 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, 1, 1, bias=False),
                                     norm_layer(inter_channels))
        self.conv2_2 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, 1, 1, bias=False),
                                     norm_layer(inter_channels))

        #self.conv2_3：使用 (1, 3) 的卷积核，只在水平方向有感受野，垂直方向不变。
        self.conv2_3 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, (1, 3), 1, (0, 1), bias=False),
                                     norm_layer(inter_channels))
        #self.conv2_4：使用 (3, 1) 的卷积核，只在垂直方向有感受野，水平方向不变。
        self.conv2_4 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, (3, 1), 1, (1, 0), bias=False),
                                     norm_layer(inter_channels))
        #self.conv2_5 和 self.conv2_6：用于融合特征并应用激活函数。
        self.conv2_5 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, 1, 1, bias=False),
                                     norm_layer(inter_channels),
                                     nn.ReLU(True))
        self.conv2_6 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, 1, 1, bias=False),
                                     norm_layer(inter_channels),
                                     nn.ReLU(True))
        #self.conv3：将两个特征图在通道维度上拼接后，通过 1x1 卷积恢复到原始通道数。
        self.conv3 = nn.Sequential(nn.Conv2d(inter_channels * 2, in_channels, 1, bias=False),
                                   norm_layer(in_channels))
        #保存上采样参数
        self._up_kwargs = up_kwargs


    def forward(self, x):
        _, _, h, w = x.size()
        #对输入特征图分别通过 conv1_1 和 conv1_2，得到两个特征图 x1 和 x2。
        x1 = self.conv1_1(x)
        x2 = self.conv1_2(x)
        #x2_1：对 x1 进行标准卷积。
        x2_1 = self.conv2_0(x1)
        #x2_2：对 x1 进行池化（尺寸为 pool_size[0]），然后卷积，上采样回原始尺寸。
        x2_2 = F.interpolate(self.conv2_1(self.pool1(x1)), (h, w), **self._up_kwargs)
        #2_3：对 x1 进行池化（尺寸为 pool_size[1]），然后卷积，上采样回原始尺寸。
        x2_3 = F.interpolate(self.conv2_2(self.pool2(x1)), (h, w), **self._up_kwargs)
        #x2_4：对 x2 在水平方向进行全局池化，卷积，上采样回原始尺寸。
        x2_4 = F.interpolate(self.conv2_3(self.pool3(x2)), (h, w), **self._up_kwargs)
        #x2_5：对 x2 在垂直方向进行全局池化，卷积，上采样回原始尺寸。
        x2_5 = F.interpolate(self.conv2_4(self.pool4(x2)), (h, w), **self._up_kwargs)
        #x1：将 x2_1、x2_2、x2_3 相加，激活后通过卷积。
        x1 = self.conv2_5(F.relu_(x2_1 + x2_2 + x2_3))
        #x2：将 x2_5 和 x2_4 相加，激活后通过卷积。
        x2 = self.conv2_6(F.relu_(x2_5 + x2_4))
        #将 x1 和 x2 在通道维度上拼接，通过 conv3 恢复通道数。
        out = self.conv3(torch.cat([x1, x2], dim=1))
        #将结果与输入 x 相加，应用激活函数，形成残差连接，得到最终输出。
        return F.relu_(x + out)

if __name__ == '__main__':
    DEVICE = 'mps' if torch.backends.mps.is_available() else 'cpu'  # 支持 MPS 则使用，否则使用 CPU
    print(f"Using device: {DEVICE}")
    block = StripPooling(
        64, (16, 8), nn.BatchNorm2d, {'mode': 'bilinear', 'align_corners': True}).to(DEVICE)

    input = torch.rand(4, 64, 64, 64).to(DEVICE)
    output = block(input)
    print(input.size(), output.size())
